package cn.bellychang.sword.finger.offer.question040;

import java.util.Arrays;

/**
 * @author ChangLiang
 * @date 2021/2/2
 */
public class Solution2 {

    public int[] getLeastNumbers(int[] arr, int k) {

        if (k == 0) {
            return new int[]{};
        }

        int index = k-1;
        getKthNumber(0,arr.length-1, arr, index);
        return Arrays.copyOfRange(arr, 0, k);
    }

    private void getKthNumber(int beginIndex, int endIndex, int[] arr, int index) {
        int partitionIndex = partition(beginIndex, endIndex, arr);
        if (partitionIndex != index) {
            if (partitionIndex < index) {
                getKthNumber(partitionIndex + 1, endIndex, arr, index);
            } else {
                getKthNumber(beginIndex, partitionIndex-1, arr, index);
            }
        }
    }

    private int partition(int beginIndex, int endIndex, int[] arr) {
        int pivot = arr[beginIndex];
        int i = beginIndex + 1;
        int j = beginIndex + 1;
        // loop invariant: [beginIndex, i-1] < pivot , 其他
        while (j <= endIndex) {
            if (arr[j] < pivot) {
                if (i != j) {
                    swap(i, j, arr);
                }
                i++;
            }
            j++;
        }
        swap(beginIndex, i - 1, arr);
        return i-1;
    }

    private void swap(int i, int j, int[] arr) {
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }
}
